package com.bruce.ai.alibaba.agent.a2a.discovery.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.cloud.ai.graph.agent.BaseAgent;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;

import jakarta.annotation.Resource;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

@RequestMapping("/root")
@RestController
public class RootAgentController {

    @Resource
    private BaseAgent rootAgent;

    private static final Logger logger = LoggerFactory.getLogger(RootAgentController.class);

    @GetMapping("/stream")
    public Flux<String> rootAgent(@RequestParam("message") String message)
            throws GraphStateException, GraphRunnerException {
        return rootAgent.stream(message).mapNotNull(output -> {
            logger.debug("stream agent invoke : `{}`", output.toString());
            if (output.isSTART() || output.isEND()) {
                return null;
            }
            if (output instanceof StreamingOutput<?> streamingOutput) {
                return streamingOutput.chunk();
            }
            return null;
        }).publishOn(Schedulers.parallel());
    }

}
